class TextGenerator:
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer

    def continue_writting(self, input_text_main_body, input_text_outline):
        context = ""
        if input_text_main_body:
            context = "\n前文内容是：" + input_text_main_body + "\n"

        if context:
            context += "根据前文内容续写，大纲是：" + input_text_outline
        else:
            context = "续写，大纲是：" + input_text_outline
    
        response, history = self.model.chat(self.tokenizer, context, history=None)
        return response
    
    # 定义头脑风暴生成函数
    # input_text_zhengwen, 
    def generate_text_tounao_baofeng(self, input_fengaoleixing):
        context =  input_fengaoleixing 

        response, history = self.model.chat(self.tokenizer, context, history=None)
        return response

    def generate_text_find_error_bug(self, input_context):
        # if input_context:
        #     #input_context
        #     a = "找错别字"
        # else:
        #     input_context = "仔细审问之后，林多文知道了罪魁祸手，原来是自己老板的女儿张雨婷。那次张雨婷根据林多文留的地址她找到了林山寨，和林多文大闹了一场，林多文收回了自己应得的6万块钱，最后不欢而散。但是在回去的路上张雨婷突然发现路过彩票店，门头上挂着庆祝的横幅。恭喜我店售卖出fl彩票一等奖。"
        
        context =  "请找出来下面文章里面的错别字:\n" + input_context

        response, history = self.model.chat(self.tokenizer, context, history=None)
        return response


    def generate_text_change(self, input_context, change_type):
        content = "请把正文的风格转换为" + change_type+":\n" + input_context

        response, history = self.model.chat(self.tokenizer, content, history=None)
        return response